package com.artislong.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.time.LocalDate;

/**
 * @author xushu
 * @version 1.0
 * @description:
 */
@RestController
@CrossOrigin
public class OpenAiController {

    private final ChatClient chatClient;
    private final ChatMemory chatMemory;

    public OpenAiController(ChatClient.Builder chatClientBuilder, ChatMemory chatMemory, EmbeddingModel embeddingModel, ToolCallbackProvider... tools) {
        this.chatClient = chatClientBuilder.defaultSystem(
                        """
                                您是“Tuling”航空公司的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                                您正在通过在线聊天系统与客户互动。
                                在提供有关预订或取消预订的信息之前，您必须始终
                                从用户处获取以下信息：预订号、客户姓名。
                                在询问用户之前，请检查消息历史记录以获取此信息。
                                在更改或退订之前，请先获取预订信息并且用户确定之后才进行更改或退订。
                                请讲中文。
                                今天的日期是 {current_date}.
                                """
                )
                .defaultAdvisors(
                        new SimpleLoggerAdvisor(),
                        // RAG
                        new QuestionAnswerAdvisor(SimpleVectorStore.builder(embeddingModel).build())
                )
                .defaultToolCallbacks(tools)
                .build();
        this.chatMemory = chatMemory;
    }

    @CrossOrigin
    @GetMapping(value = "/ai/generateStreamAsString", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generateStreamAsString(@RequestParam(value = "message", defaultValue = "讲个笑话") String message, @RequestParam(value = "conversationId") String conversationId) {
        Flux<String> content = this.chatClient.prompt()
                .user(message)
                .system(promptSystemSpec -> promptSystemSpec.param("current_date", LocalDate.now().toString()))
                .advisors(PromptChatMemoryAdvisor.builder(chatMemory).conversationId(conversationId).scheduler(Schedulers.immediate()).build())
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, 100))
                .stream()
                .content();
        return content.concatWith(Flux.just("[complete]"));
    }

}
